/*************************************************************************
	> File Name: ./include/PageMap.h
	> Author: zq
	> Mail: zhouquan1511@163.com 
	> Created Time: Sun 30 Apr 2023 10:18:38 AM CST
 ************************************************************************/
#pragma once

#include "ObjectPool.h"
#include "Common.h"

/*单层基数树:32位平台
 * BITS为当前平台的页面最大数量对应的bit位
 * 例如32位平台下2^32/2^13,即32-13=19
*/
template<int BITS>
class TCMalloc_PageMap1 {
	private:
		void **_array;
		static const int LENGTH = 1 << BITS;
		
	public:
		typedef uintptr_t Number;
		explicit TCMalloc_PageMap1() {
			size_t size = sizeof(void*) << BITS;
			size_t alignSize = SizeClass::_RoundUp(size, 1 << PAGE_SHIFT);
			_array = (void**)SystemAlloc(alignSize >> PAGE_SHIFT);
			memset(_array, 0, size);
		}
		
		void* get(Number k) const {
			if ((k >> BITS) > 0) {
				return nullptr;
			}
			return _array[k];
		}
		
		void set(Number k, void *v) {
			assert((k >> BITS) == 0);
			_array[k] = v;
		}	
}

//两层基数树 -> 32位平台
template<int BITS>
class TCMalloc_PageMap2 {
	private:
		static const int ROOT_BITS = 5; //第一层对应页号的前5位
		static const int ROOT_LENGTH = 1 << ROOT_BITS; //第一层元素个数
		static const int LEAF_BITS = BITS - ROOT_BITS; //第二层对应页号的其余位
		static const int LEAF_LENGTH = 1 << LEAF_BITS; //第二层元素个数

		struct Leaf {
			void *values[LEAF_LENGTH];
		};
		Leaf* _root[ROOT_LENGTH];

	public:
		typedef uintptr_t Number;
		explicit TCMalloc_PageMap2() {
			memset(_root, 0, sizeof(_root));
			PreallocateMoreMemory(); //提前开辟第二层空间
		}
		
		void* get(Number k) const {
			const Number i1 = k >> LEAF_BITS;
			const Number i2 = k & (LEAF_LENGTH - 1);
			if ((k >> BITS) > 0 || _root[i1] == nullptr) {
				return nullptr;
			}
			return _root[i1]->values[i2];
		}

		void set(Number k, void *v) {
			const Number i1 = k >> LEAF_BITS;
			const Number i2 = k & (LEAF_LENGTH - 1);
			assert(i1 < ROOT_LENGTH);
			if (_root[i1] == nullptr) {
				Ensure(k, 1);
			}
			_root[i1]->values[i2] = v; 
		}

		//确保[start, start+n-1]页号的空间开辟好
		bool Ensure(Number start, size_t n) {
			for (Number k = start; k <= start + n - 1;) {
				const Number i1 = k >> LEAF_BITS;
				if (i1 >= ROOT_LENGTH) { //页号超出范围
					return false;
				}
				if (_root[i1] == nullptr) {
					static ObjectPool<Leaf> leafPool;
					Leaf *leaf = (Leaf*)leafPool.New();
					_root[i1] = leaf;
				}
				k = ((k >> LEAF_BITS) + 1) << LEAF_BITS;
			}
			return true;
		}

		void PreallocateMoreMemory() {
			Ensure(0, 1 << BITS); //确保第二层的空间都开辟好
		}
}

//实现此部分代码时，全程将k想象为二进制表现形式
//其位数结构为: BITS = INTERIOR_BITS + INTERIOR_BITS + LEAF_BITS
/* 三层基数树->64位平台
 * 只有当建立某一页号的映射时，再开辟对应的数组空间
 * 没有建立映射的页号不需要开辟其对应的数组空间
 * 可节省大量内存
 */
template<int BITS>
class TCMalloc_PageMap3 {
	private:
		static const int INTERIOR_BITS = (BITS + 2) / 3; //一、二层
		static const int INTERIOR_LENGTH = 1 << INTERIOR_BITS; 
		static const int LEAF_BITS = BITS - INTERIOR_BITS * 2; //三层
		static const int LEAF_LENGTH = 1 << LEAF_BITS;
		
		struct Node {
			Node* ptrs[INTERIOR_LENGTH];
		};
		struct Leaf {
			void* values[LEAF_LENGTH];
		};
		Node* _root;

		Node* NewNode() {
			static ObjectPool<Node> nodePool;
			Node *result = (Node*)nodePool.New();
			if (result != nullptr) {
				memset(result, 0, sizeof(*result));
			}
			return result;
		}

	public:
		typedef uintptr_t Number;
		explicit TCMalloc_PageMap3() {
			_root = NewNode();
		}

		void* get(Number k) {
			const Number i1 = k >> (INTERIOR_BITS + LEAF_BITS);
			const Number i2 = (k >> LEAF_BITS) & (INTERIOR_LENGTH - 1);
			const Number i3 = k & (LEAF_LENGTH - 1);
			if ((k >> BITS) > 0 || _root->ptrs[i1] == nullptr
								|| _root->ptrs[i1]->ptrs[i2] == nullptr) {
				return nullptr;
			}
			return reinterpret_cast<Leaf*>(_root->ptrs[i1]->ptrs[i2])->values[i3];
		}
		
		void set(Number k, void *v) {
			assert((k >> BITS) == 0);	
			const Number i1 = k >> (INTERIOR_BITS + LEAF_BITS);
			const Number i2 = (k >> LEAF_BITS) & (INTERIOR_LENGTH - 1);
			const Number i3 = k & (LEAF_LENGTH - 1);
			Ensure(k, 1);
			reinterpret_cast<Leaf*>(_root->ptrs[i1]->ptrs[i2])->values[k] = v;
		}

		//确保[start, start+n-1]页号对应的映射空间已开辟好
		bool Ensure(Number start, size_t n) {
			for (Number k = start; k <= start + n - 1;) {
				const Number i1 = k >> (INTERIOR_BITS + LEAF_BITS);
				const Number i2 = (k >> LEAF_BITS) & (INTERIOR_LENGTH - 1);
				if (i1 >= INTERIOR_LENGTH || i2 >= INTERIOR_LENGTH) {
					return false;
				}
				//构建第一层
				if (_root->ptrs[i1] == nullptr) {
					Node *node = NewNode();
					if (node == nullptr)
						return false;
					_root[i1] = node;
				}
				//构建第二层
				if (_root->ptrs[i1]->ptrs[i2] == nullptr) {
					static ObjectPool<Leaf> leafPool;
					Leaf *leaf = leafPool.New();
					if (leaf == nullptr)
						return false;
					memset(leaf, 0, sizeof(*leaf));
					_root->ptrs[i1]->ptrs[i2] = reinterpret_cast<Node*>(leaf);
				}
				k = ((k >> LEAF_BITS) + 1) << LEAF_BITS;
			}
			return true;
		}
}
